package com.demo.newLSTM;

import org.deeplearning4j.datasets.iterator.impl.ListDataSetIterator;
import org.deeplearning4j.nn.conf.ComputationGraphConfiguration;
import org.deeplearning4j.nn.conf.MultiLayerConfiguration;
import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
import org.deeplearning4j.nn.conf.layers.DenseLayer;
import org.deeplearning4j.nn.conf.layers.LSTM;
import org.deeplearning4j.nn.conf.layers.RnnOutputLayer;
import org.deeplearning4j.nn.graph.ComputationGraph;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.nn.weights.WeightInit;
import org.nd4j.linalg.activations.Activation;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.DataSet;
import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.learning.config.Adam;
import org.nd4j.linalg.lossfunctions.LossFunctions;

import java.util.ArrayList;
import java.util.List;

public class test {
    public static void main(String[] args) {

        // 将输入数据和标签打包成DataSet对象
        DataSet dataSet = setINDArray();

        // 将数据集拆分为训练集和测试集
        List<DataSet> list = dataSet.asList();
        int trainSize = (int) (list.size() * 0.8); // 使用80%的数据作为训练集
        List<DataSet> trainList = list.subList(0, trainSize);

        DataSetIterator trainIterator = new ListDataSetIterator(trainList);

        // 构建LSTM网络模型
       /* ComputationGraphConfiguration conf = new NeuralNetConfiguration.Builder()
                .seed(123)
                .weightInit(WeightInit.XAVIER)
                .updater(new Adam(1e-3))
                .graphBuilder()
                .addInputs("input")
                .addLayer("LSTM", new LSTM.Builder().nIn(numInputs).nOut(10).build(), "input")
                .addLayer("dense", new DenseLayer.Builder().nIn(10).nOut(1).activation("identity").build(), "LSTM")
                .setOutputs("dense")
                .build();
        ComputationGraph net = new ComputationGraph(conf);
        net.init();*/

        MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder()
                .seed(123)
                .weightInit(WeightInit.XAVIER)
                .updater(new Adam(1e-3))
                .list()
                .layer(0, new LSTM.Builder()
                        .nIn(3)
                        .nOut(10)
                        .activation(Activation.TANH)
                        .gateActivationFunction(Activation.SIGMOID)
                        .dropOut(0.5)
                        .build())
                .layer(new RnnOutputLayer.Builder(LossFunctions.LossFunction.MSE)
                        .activation(Activation.IDENTITY)
                        .nIn(10)
                        .nOut(1)
                        .build())
                .build();
        MultiLayerNetwork net = new MultiLayerNetwork(conf);
        net.init();

        // 训练模型
        for (int i = 0; i < 100; i++) {
            //net.fit(trainIterator);
            //System.out.println("Epoch " + i + " 完成");

            net.fit(trainIterator);
            trainIterator.reset();
        }

        /*while (trainIterator.hasNext()) {
            DataSet testData = trainIterator.next();
            INDArray output = net.output(testData.getFeatures())[…[omitted]

            // 输出预测结果
            System.out.println("预测负荷值：" + output);
        }*/

        double[] load = {10.0, 12.0, 15.0, 20.0, 18.0, 14.0, 12.0, 10.0, 9.0, 8.0, 9.0, 11.0, 13.0, 15.0, 18.0, 20.0};
        double[] temperature = {20.0, 22.0, 25.0, 28.0, 30.0, 28.0, 25.0, 22.0, 20.0, 18.0, 16.0, 18.0, 20.0, 22.0, 25.0, 28.0};
        double[] humidity = {50.0, 55.0, 60.0, 65.0, 70.0, 75.0, 80.0, 85.0, 90.0, 85.0, 80.0, 75.0, 70.0, 65.0, 60.0, 55.0};

        int numInputs = 3;
        //使用训练好的模型进行预测
        INDArray inputForPrediction = Nd4j.zeros(1, numInputs);
        inputForPrediction.putScalar(new int[] {0, 0}, load[load.length - numInputs]);
        inputForPrediction.putScalar(new int[] {0, 1}, temperature[temperature.length - numInputs]);
        inputForPrediction.putScalar(new int[] {0, 2}, humidity[humidity.length - numInputs]);
        for (int i = 1; i <= 4; i++) {
            INDArray outputForPrediction = net.output(inputForPrediction);
            double predictedValue = outputForPrediction.getDouble(0);
            System.out.println("预测值：" + predictedValue);
            inputForPrediction.putScalar(new int[] {0, 0}, predictedValue);
            inputForPrediction.putScalar(new int[] {0, 1}, temperature[temperature.length - numInputs + i]);
            inputForPrediction.putScalar(new int[] {0, 2}, humidity[humidity.length - numInputs + i]);
        }

        // 在测试集上评估模型性能
        /*
        double mseSum = 0.0;
        int count = 0;

        while (testIterator.hasNext()) {
            DataSet testData = testIterator.next();
            INDArray output = net.output(testData.getFeatures())[…[omitted]

        // 使用训练好的模型进行预测
        INDArray inputForPrediction = Nd4j.zeros(1, numInputs);
        inputForPrediction.putScalar(new int[] {0, 0}, load[load.length - numInputs]);
        inputForPrediction.putScalar(new int[] {0, 1}, temperature[temperature.length - numInputs]);
        inputForPrediction.putScalar(new int[] {0, 2}, humidity[humidity.length - numInputs]);
        for (int i = 1; i <= 4; i++) {
            INDArray outputForPrediction = net.output(inputForPrediction);
            double predictedValue = outputForPrediction.getDouble(0);
            System.out.println("预测值：" + predictedValue);
            inputForPrediction.putScalar(new int[] {0, 0}, predictedValue);
            inputForPrediction.putScalar(new int[] {0, 1}, temperature[temperature.length - numInputs + i]);
            inputForPrediction.putScalar(new int[] {0, 2}, humidity[humidity.length - numInputs + i]);
        }*/
    }


    public static DataSet setINDArray(){
        // 假设已经从数据库或其他数据源中获取到了负荷数据、温度数据和湿度数据等
        double[] load = {10.0, 12.0, 15.0, 20.0, 18.0, 14.0, 12.0, 10.0, 9.0, 8.0, 9.0, 11.0, 13.0, 15.0, 18.0, 20.0};
        double[] temperature = {20.0, 22.0, 25.0, 28.0, 30.0, 28.0, 25.0, 22.0, 20.0, 18.0, 16.0, 18.0, 20.0, 22.0, 25.0, 28.0};
        double[] humidity = {50.0, 55.0, 60.0, 65.0, 70.0, 75.0, 80.0, 85.0, 90.0, 85.0, 80.0, 75.0, 70.0, 65.0, 60.0, 55.0};

        // 准备输入数据
        int numExamples = load.length - 4 - 4; // 预测4个小时，因此减去前4个数据和后4个数据
        int numInputs = 3; // 加入温度和湿度，输入维度为3
        INDArray input = Nd4j.zeros(numExamples, numInputs);
        INDArray labels = Nd4j.zeros(numExamples, 1);
        for (int i = 0; i < numExamples; i++) {
            input.putScalar(new int[] {i, 0}, load[i]);
            input.putScalar(new int[] {i, 1}, temperature[i]);
            input.putScalar(new int[] {i, 2}, humidity[i]);
            labels.putScalar(new int[] {i, 0}, load[i + 4]);
        }

        // 将输入数据和标签打包成DataSet对象
        DataSet dataSet = new DataSet(input, labels);
        return dataSet;
    }
}
